import numpy
import matplotlib.pyplot as plt
import pandas
import json
import os
import utm
from math import pi, floor

# 角度限制, 输入弧度, 输出-pi~pi
def rad_limit(rad: float) -> float:
	if rad >= 180:
		rad = rad - 360
	if rad < -180:
		rad = rad + 360
	return rad

# json to dict?
def json_decode(df_data: pandas.DataFrame):
    columns_name = []
    for key in json.loads(df_data.iloc[0]):
        columns_name.append(key)
    
    temp = pandas.DataFrame(columns=columns_name)
    for i in range(len(df_data)):
        temp_dict = json.loads(df_data.iloc[i])
        temp_series = pandas.Series(temp_dict).to_frame()
        temp = pandas.concat([temp, temp_series.T])
    return temp

def cal_psi_error(data_df: pandas.DataFrame):
    data_df['psi_error'] = data_df['psi'] - data_df['exp_psi']
    for i in range(len(data_df['psi_error'])):
        data_df['psi_error'].iloc[i] = rad_limit(data_df['psi_error'].iloc[i])
    return data_df

if __name__ == '__main__':
    HOME = os.getcwd()
    
    exp_line = {'lon':[124.16168462, 124.18153454, 124.22066772],
                'lat':[29.64351389, 29.650637969, 29.645361697]}
    exp_line_df = pandas.DataFrame(exp_line)

    data_df = pandas.read_csv(f'{HOME}/data_deal/2024-10-23status.csv')
    # print(data_df.columns)

    pid_data = pandas.DataFrame()
    mpc_data = pandas.DataFrame()
    td3_data = pandas.DataFrame()
    sacn_data = pandas.DataFrame()

    pid_data = data_df["drl"][22:22+1180]
    pid_data_df = json_decode(pid_data)
    pid_data_df["rud"] = json_decode(data_df["control"][22:22+1180])["rudl"]
    pid_data_df = cal_psi_error(pid_data_df)
    print(pid_data_df.columns)

    mpc_data = data_df["drl"][1440:1440+1180]
    mpc_data_df = json_decode(mpc_data)
    mpc_data_df = cal_psi_error(mpc_data_df)
    
    td3_data = data_df["drl"][3155:3155+1180]
    td3_data_df = json_decode(td3_data)
    td3_data_df = cal_psi_error(td3_data_df)
    
    sacn_data = data_df["drl"][4992:4992+1180]
    sacn_data_df = json_decode(sacn_data)
    sacn_data_df = cal_psi_error(sacn_data_df)
    
    plt.rcParams['font.size']=18
    fig, ax = plt.subplots(ncols=2, nrows=2, figsize=(16,9))
    ax[0,0].plot(pid_data_df['psi_error'].to_numpy(), label='PID')
    ax[0,0].legend()
    ax[0,0].set_xlabel('Time / s')
    ax[0,0].set_ylabel('Heading angle error / $^o$')
    ax[0,0].set_ylim([-40, 25])

    ax[0,1].plot(mpc_data_df['psi_error'].to_numpy(), label='NMPC')
    ax[0,1].legend()
    ax[0,1].set_xlabel('Time / s')
    ax[0,1].set_ylabel('Heading angle error / $^o$')
    ax[0,1].set_ylim([-40, 25])

    ax[1,0].plot(td3_data_df['psi_error'].to_numpy(), label='TD3')
    ax[1,0].legend()
    ax[1,0].set_xlabel('Time / s')
    ax[1,0].set_ylabel('Heading angle error / $^o$')
    ax[1,0].set_ylim([-40, 25])

    ax[1,1].plot(sacn_data_df['psi_error'].to_numpy(), label='SAC-N')
    ax[1,1].legend()
    ax[1,1].set_xlabel('Time / s')
    ax[1,1].set_ylabel('Heading angle error / $^o$')
    ax[1,1].set_ylim([-40, 25])

    plt.savefig(f'{HOME}/data_deal/plot_psi_error.eps', bbox_inches='tight', dpi=300)
    
    print('pid mae is', numpy.abs(pid_data_df['psi_error'].to_numpy()).mean())
    print('mpc mae is', numpy.abs(mpc_data_df['psi_error'].to_numpy()).mean())
    print('td3 mae is', numpy.abs(td3_data_df['psi_error'].to_numpy()).mean())
    print('sac mae is', numpy.abs(sacn_data_df['psi_error'].to_numpy()).mean())
    
    plt.show()